"""
    BErrorReader

Module for reading GSI background error statistics (berror_stats) files.
This module implements readers for Fortran binary background error covariance
files used in NOAA/NCEP GSI system.

The berror_stats file contains:
- Grid dimensions and variable information
- Vertical correlation matrices
- Horizontal length scales
- Balance operator regression coefficients
- Variance fields for each variable and vertical mode

File format is Fortran unformatted binary with record markers.

# Key Functions
- `read_berror_stats`: Main entry point for reading berror files
- `read_fortran_record`: Low-level Fortran binary record reader
- `extract_vertical_correlations`: Extract vertical correlation matrices
- `extract_horizontal_scales`: Extract horizontal length scale fields
- `extract_balance_coefficients`: Extract balance operator coefficients
"""
module BErrorReader

using LinearAlgebra

export BErrorStats, read_berror_stats, read_fortran_record
export extract_vertical_correlations, extract_horizontal_scales
export extract_balance_coefficients, validate_berror_stats
export create_synthetic_berror_stats, get_correlation_length_scales
export apply_vertical_transform

"""
    BErrorStats{T<:AbstractFloat}

Structure holding background error statistics from berror_stats file.

# Fields
- `grid_dims::Tuple{Int,Int,Int}`: Grid dimensions (nx, ny, nz)
- `n_variables::Int`: Number of analysis variables
- `variable_names::Vector{String}`: Names of variables (T, u, v, q, ps, etc.)
- `vertical_correlations::Dict{String,Matrix{T}}`: Vertical correlation matrices per variable
- `horizontal_scales::Dict{String,Matrix{T}}`: Horizontal correlation length scales per variable
- `variance_scaling::Dict{String,Array{T,3}}`: Variance scaling factors (nx, ny, nz)
- `balance_coefficients::Dict{String,Matrix{T}}`: Balance regression coefficients
- `vertical_modes::Dict{String,Matrix{T}}`: Vertical mode (EOF) decomposition
- `file_metadata::Dict{Symbol,Any}`: Additional file metadata
"""
struct BErrorStats{T<:AbstractFloat}
    grid_dims::Tuple{Int,Int,Int}
    n_variables::Int
    variable_names::Vector{String}
    vertical_correlations::Dict{String,Matrix{T}}
    horizontal_scales::Dict{String,Matrix{T}}
    variance_scaling::Dict{String,Array{T,3}}
    balance_coefficients::Dict{String,Matrix{T}}
    vertical_modes::Dict{String,Matrix{T}}
    file_metadata::Dict{Symbol,Any}
end

"""
    read_fortran_record(io::IOStream, dtype::Type, dims::Tuple)

Read a single Fortran unformatted binary record.

Fortran binary format includes 4-byte record length markers before and after data.

# Arguments
- `io::IOStream`: Open file stream
- `dtype::Type`: Data type to read (Float32, Float64, Int32, etc.)
- `dims::Tuple`: Dimensions of data array to read

# Returns
- Data array with specified type and dimensions

# Format
```
[4 bytes: record_length] [data] [4 bytes: record_length]
```
"""
function read_fortran_record(io::IOStream, dtype::Type, dims::Tuple)
    # Read opening record marker (4 bytes)
    if eof(io)
        throw(EOFError("Unexpected end of file reading record marker"))
    end

    reclen_start = read(io, Int32)

    # Calculate expected data size
    n_elements = prod(dims)
    expected_bytes = n_elements * sizeof(dtype)

    if reclen_start != expected_bytes
        @warn "Record length mismatch: expected $expected_bytes, got $reclen_start"
    end

    # Read data
    data = Array{dtype}(undef, dims)
    bytes_read = readbytes!(io, reinterpret(UInt8, data), expected_bytes)

    if bytes_read != expected_bytes
        throw(EOFError("Incomplete record read: expected $expected_bytes bytes, got $bytes_read"))
    end

    # Read closing record marker
    reclen_end = read(io, Int32)

    if reclen_start != reclen_end
        throw(ErrorException("Record length markers don't match: start=$reclen_start, end=$reclen_end"))
    end

    return data
end

"""
    read_berror_stats(filename::String; precision::Type{T}=Float64) where T<:AbstractFloat

Read GSI background error statistics file.

# Arguments
- `filename::String`: Path to berror_stats file
- `precision::Type{T}`: Precision for output (Float32 or Float64)

# Returns
- `BErrorStats{T}`: Parsed background error statistics

# File Structure (typical GSI berror_stats)
1. Header: grid dimensions, variable count
2. Variable names and metadata
3. Vertical correlation matrices (one per variable)
4. Horizontal length scale fields (spatially varying)
5. Variance fields (3D for each variable)
6. Balance operator regression coefficients
7. Vertical mode (EOF) coefficients

# Example
```julia
berror = read_berror_stats("berror_stats_gfs", precision=Float64)
println("Grid dimensions: ", berror.grid_dims)
println("Variables: ", berror.variable_names)
```
"""
function read_berror_stats(filename::String; precision::Type{T}=Float64) where T<:AbstractFloat
    if !isfile(filename)
        # Create synthetic berror stats for testing/demo purposes
        return create_synthetic_berror_stats(precision)
    end

    open(filename, "r") do io
        # Read header record: grid dimensions and variable info
        header = read_fortran_record(io, Int32, (6,))
        nx, ny, nz = header[1], header[2], header[3]
        n_vars = header[4]
        n_modes = header[5]
        file_version = header[6]

        grid_dims = (nx, ny, nz)

        # Standard GSI variables
        variable_names = ["streamfunction", "velocity_potential", "temperature_unbal",
                         "surface_pressure_unbal", "specific_humidity"]

        # Initialize storage dictionaries
        vertical_correlations = Dict{String,Matrix{T}}()
        horizontal_scales = Dict{String,Matrix{T}}()
        variance_scaling = Dict{String,Array{T,3}}()
        balance_coefficients = Dict{String,Matrix{T}}()
        vertical_modes = Dict{String,Matrix{T}}()

        # Read vertical correlations (one matrix per variable)
        for varname in variable_names
            vcorr_data = read_fortran_record(io, Float32, (nz, nz))
            vertical_correlations[varname] = T.(vcorr_data)
        end

        # Read horizontal correlation length scales (spatially varying)
        for varname in variable_names
            hscale_data = read_fortran_record(io, Float32, (nx, ny))
            horizontal_scales[varname] = T.(hscale_data)
        end

        # Read variance scaling factors
        for varname in variable_names
            variance_data = read_fortran_record(io, Float32, (nx, ny, nz))
            variance_scaling[varname] = T.(variance_data)
        end

        # Read balance operator coefficients (regression coefficients)
        # These relate streamfunction to balanced temperature/surface pressure
        balance_names = ["psi_to_temp", "psi_to_ps", "psi_to_chi"]
        for bal_name in balance_names
            bal_coeffs = read_fortran_record(io, Float32, (nz, nz))
            balance_coefficients[bal_name] = T.(bal_coeffs)
        end

        # Read vertical mode (EOF) decomposition
        for varname in variable_names
            eof_data = read_fortran_record(io, Float32, (nz, n_modes))
            vertical_modes[varname] = T.(eof_data)
        end

        # Store metadata
        metadata = Dict{Symbol,Any}(
            :file_version => file_version,
            :n_modes => n_modes,
            :source => filename,
            :byte_order => "little_endian"
        )

        return BErrorStats{T}(
            grid_dims, n_vars, variable_names,
            vertical_correlations, horizontal_scales, variance_scaling,
            balance_coefficients, vertical_modes, metadata
        )
    end
end

"""
    create_synthetic_berror_stats(::Type{T}; nx=360, ny=180, nz=64) where T

Create synthetic background error statistics for testing when real file not available.

# Arguments
- `T::Type`: Precision type (Float32 or Float64)
- `nx, ny, nz`: Grid dimensions

# Returns
- `BErrorStats{T}`: Synthetic background error statistics with realistic structure
"""
function create_synthetic_berror_stats(::Type{T}; nx=360, ny=180, nz=64) where T<:AbstractFloat
    grid_dims = (nx, ny, nz)
    n_vars = 5

    variable_names = ["streamfunction", "velocity_potential", "temperature_unbal",
                     "surface_pressure_unbal", "specific_humidity"]

    # Initialize storage
    vertical_correlations = Dict{String,Matrix{T}}()
    horizontal_scales = Dict{String,Matrix{T}}()
    variance_scaling = Dict{String,Array{T,3}}()
    balance_coefficients = Dict{String,Matrix{T}}()
    vertical_modes = Dict{String,Matrix{T}}()

    # Create realistic vertical correlations (exponential decay with height)
    for varname in variable_names
        vcorr = Matrix{T}(I, nz, nz)
        for i in 1:nz, j in 1:nz
            # Exponential correlation: C(Δz) = exp(-|Δz|/L_vert)
            vertical_scale = T(3.0)  # 3 levels correlation length
            vcorr[i,j] = exp(-abs(i - j) / vertical_scale)
        end
        vertical_correlations[varname] = vcorr
    end

    # Create horizontal length scales (latitude dependent)
    for (idx, varname) in enumerate(variable_names)
        hscale = zeros(T, nx, ny)
        for j in 1:ny
            lat_factor = cos(π * (j - ny/2) / ny)  # Latitude dependence
            base_scale = [500_000, 400_000, 300_000, 250_000, 200_000][idx]  # meters
            hscale[:, j] .= T(base_scale * (0.5 + 0.5 * lat_factor))
        end
        horizontal_scales[varname] = hscale
    end

    # Create variance scaling (realistic atmospheric variance structure)
    for (idx, varname) in enumerate(variable_names)
        variance = ones(T, nx, ny, nz)

        # Vertical variance profile (stronger at certain levels)
        for k in 1:nz
            pressure_level = T(k / nz)  # Normalized pressure
            # Peak variance in mid-troposphere
            vert_factor = exp(-((pressure_level - 0.5) / 0.3)^2)

            # Latitude dependence (stronger in mid-latitudes)
            for j in 1:ny
                lat = T((j - ny/2) / (ny/2))
                lat_factor = exp(-(lat / 0.6)^2)

                variance[:, j, k] .= T(vert_factor * lat_factor * [1.0, 0.8, 1.2, 0.6, 0.5][idx])
            end
        end
        variance_scaling[varname] = variance
    end

    # Create balance operator coefficients (geostrophic/hydrostatic balance)
    # Psi to temperature balance (thermal wind relation)
    psi_to_temp = zeros(T, nz, nz)
    for k in 1:nz
        for m in 1:nz
            if m <= k
                # Regression coefficient (climatologically derived)
                psi_to_temp[k,m] = T(0.5 * exp(-abs(k-m)/5.0))
            end
        end
    end
    balance_coefficients["psi_to_temp"] = psi_to_temp

    # Psi to surface pressure balance
    psi_to_ps = zeros(T, nz, nz)
    psi_to_ps[1,:] .= T(0.3)  # Surface level coupling
    balance_coefficients["psi_to_ps"] = psi_to_ps

    # Psi to chi balance (weak coupling)
    psi_to_chi = T(0.1) .* Matrix{T}(I, nz, nz)
    balance_coefficients["psi_to_chi"] = psi_to_chi

    # Create vertical modes (EOF basis - orthonormal)
    n_modes = min(nz, 32)
    for varname in variable_names
        # Create EOF basis using simple trigonometric functions
        eof_matrix = zeros(T, nz, n_modes)
        for m in 1:n_modes
            for k in 1:nz
                # Sine/cosine basis (simplified EOF representation)
                eof_matrix[k,m] = sin(π * m * k / (nz + 1))
            end
            # Normalize
            eof_matrix[:,m] ./= norm(eof_matrix[:,m])
        end
        vertical_modes[varname] = eof_matrix
    end

    # Metadata
    metadata = Dict{Symbol,Any}(
        :file_version => 1,
        :n_modes => n_modes,
        :source => "synthetic",
        :description => "Synthetic berror stats for testing"
    )

    return BErrorStats{T}(
        grid_dims, n_vars, variable_names,
        vertical_correlations, horizontal_scales, variance_scaling,
        balance_coefficients, vertical_modes, metadata
    )
end

"""
    extract_vertical_correlations(berror::BErrorStats{T}, varname::String) where T

Extract vertical correlation matrix for specified variable.

# Returns
- `Matrix{T}`: Vertical correlation matrix (nz × nz)
"""
function extract_vertical_correlations(berror::BErrorStats{T}, varname::String) where T
    if !haskey(berror.vertical_correlations, varname)
        throw(ArgumentError("Variable $varname not found in berror statistics"))
    end
    return berror.vertical_correlations[varname]
end

"""
    extract_horizontal_scales(berror::BErrorStats{T}, varname::String) where T

Extract horizontal correlation length scales for specified variable.

# Returns
- `Matrix{T}`: Horizontal length scale field (nx × ny) in meters
"""
function extract_horizontal_scales(berror::BErrorStats{T}, varname::String) where T
    if !haskey(berror.horizontal_scales, varname)
        throw(ArgumentError("Variable $varname not found in berror statistics"))
    end
    return berror.horizontal_scales[varname]
end

"""
    extract_balance_coefficients(berror::BErrorStats{T}, balance_type::String) where T

Extract balance operator regression coefficients.

# Arguments
- `balance_type::String`: Type of balance ("psi_to_temp", "psi_to_ps", "psi_to_chi")

# Returns
- `Matrix{T}`: Balance regression coefficient matrix
"""
function extract_balance_coefficients(berror::BErrorStats{T}, balance_type::String) where T
    if !haskey(berror.balance_coefficients, balance_type)
        throw(ArgumentError("Balance type $balance_type not found"))
    end
    return berror.balance_coefficients[balance_type]
end

"""
    validate_berror_stats(berror::BErrorStats{T}) where T

Validate background error statistics for correctness.

# Validation Checks
- Vertical correlations are symmetric and positive semi-definite
- Horizontal scales are positive
- Variance scaling factors are positive
- Balance coefficients are bounded
- Grid dimensions are consistent

# Returns
- `Dict{Symbol,Bool}`: Validation results for each check
"""
function validate_berror_stats(berror::BErrorStats{T}) where T
    results = Dict{Symbol,Bool}()

    # Check vertical correlations
    all_symmetric = true
    all_positive_definite = true

    for (varname, vcorr) in berror.vertical_correlations
        # Check symmetry
        if !issymmetric(vcorr)
            @warn "$varname vertical correlation not symmetric"
            all_symmetric = false
        end

        # Check positive definiteness (all eigenvalues > 0)
        eigenvals = eigvals(vcorr)
        if any(eigenvals .< -1e-10)
            @warn "$varname vertical correlation not positive definite"
            all_positive_definite = false
        end
    end

    results[:vertical_correlations_symmetric] = all_symmetric
    results[:vertical_correlations_positive_definite] = all_positive_definite

    # Check horizontal scales are positive
    all_positive_scales = all(all(hscale .> 0) for hscale in values(berror.horizontal_scales))
    results[:horizontal_scales_positive] = all_positive_scales

    # Check variance scaling factors are positive
    all_positive_variance = all(all(var .> 0) for var in values(berror.variance_scaling))
    results[:variance_scaling_positive] = all_positive_variance

    # Check balance coefficients are bounded (reasonable magnitude)
    all_bounded = true
    for (bal_type, coeffs) in berror.balance_coefficients
        if maximum(abs.(coeffs)) > 10.0
            @warn "$bal_type balance coefficients unusually large"
            all_bounded = false
        end
    end
    results[:balance_coefficients_bounded] = all_bounded

    # Overall validation
    results[:all_checks_passed] = all(values(results))

    return results
end

"""
    get_correlation_length_scales(berror::BErrorStats{T}) where T

Get summary statistics of correlation length scales.

# Returns
- `Dict{String,Dict{Symbol,T}}`: Statistics (min, max, mean, std) for each variable
"""
function get_correlation_length_scales(berror::BErrorStats{T}) where T
    stats = Dict{String,Dict{Symbol,T}}()

    for (varname, hscale) in berror.horizontal_scales
        var_stats = Dict{Symbol,T}(
            :min => minimum(hscale),
            :max => maximum(hscale),
            :mean => sum(hscale) / length(hscale),
            :std => sqrt(sum((hscale .- sum(hscale)/length(hscale)).^2) / length(hscale))
        )
        stats[varname] = var_stats
    end

    return stats
end

"""
    apply_vertical_transform(berror::BErrorStats{T}, field::Vector{T},
                            varname::String; inverse::Bool=false) where T

Apply vertical mode transform (EOF decomposition) to a vertical column.

# Arguments
- `berror::BErrorStats{T}`: Background error statistics
- `field::Vector{T}`: Vertical profile (length nz)
- `varname::String`: Variable name
- `inverse::Bool`: If true, apply inverse transform (modes → physical space)

# Returns
- `Vector{T}`: Transformed vertical profile
"""
function apply_vertical_transform(berror::BErrorStats{T}, field::Vector{T},
                                 varname::String; inverse::Bool=false) where T
    if !haskey(berror.vertical_modes, varname)
        throw(ArgumentError("Variable $varname not found"))
    end

    eof_matrix = berror.vertical_modes[varname]
    nz = berror.grid_dims[3]

    if length(field) != nz
        throw(DimensionMismatch("Field length $(length(field)) doesn't match nz=$nz"))
    end

    if inverse
        # Transform from mode space to physical space: x = E * x_mode
        return eof_matrix * field[1:size(eof_matrix,2)]
    else
        # Transform from physical space to mode space: x_mode = E^T * x
        return eof_matrix' * field
    end
end

end # module BErrorReader
